// This is the main DLL file.

#include "stdafx.h"

#include "SSLLib.h"

#include <msclr/marshal.h>
using namespace msclr::interop;

std::string vformat(const char* fmt, va_list ap)
{
	const int INIT_BUF_SIZE = 100;

	int n, size = INIT_BUF_SIZE;
	char *p, *np;

	if ((p = (char*)malloc (size)) == NULL)
		throw new std::bad_alloc();

	while (1)
	{
		n = vsnprintf (p, size, fmt, ap);
		if (n > -1 && n < size)
			return std::string(p);
		if (n > -1)
			size = n+1;
		else
			size *= 2;
		if ((np = (char*)realloc (p, size)) == NULL)
		{
			free(p);
			throw new std::bad_alloc();
		}
		else
			p = np;
	}
}

std::string format(const char* fmt, ...)
{
	va_list ap;
	va_start(ap, fmt);
	std::string buf = vformat(fmt, ap);
	va_end(ap);
	return buf;
}

class Buffers : public SecBufferDesc
{
	protected:
	SecBuffer* bufs;
	int nb;
	SecurityFunctionTable* sf;

	public:
	Buffers(int nBufs)
	{
		nb = nBufs;
		bufs = new SecBuffer[nBufs];
		cBuffers = nb;
		pBuffers = bufs;
		ulVersion = SECBUFFER_VERSION;
	}

	~Buffers() { delete[] bufs; }

	void SetBuffer(int i, unsigned long buftype, void* buf, int size)
	{
		bufs[i].BufferType = buftype;
		bufs[i].cbBuffer = size;
		bufs[i].pvBuffer = buf;
	}

	void SetBufferToken(int i, void* buf, int size) { SetBuffer(i, SECBUFFER_TOKEN, buf, size); }
	void SetBufferStreamHeader(int i, void* buf, int size) { SetBuffer(i, SECBUFFER_STREAM_HEADER, buf, size); }
	void SetBufferStreamTrailer(int i, void* buf, int size) { SetBuffer(i, SECBUFFER_STREAM_TRAILER, buf, size); }
	void SetBufferData(int i, void* buf, int size) { SetBuffer(i, SECBUFFER_DATA, buf, size); }
	void SetBufferEmpty(int i) { SetBuffer(i, SECBUFFER_EMPTY, NULL, 0); }

	SecBuffer* First(int type) { for (int i = 0; i < nb; i ++) if (bufs[i].BufferType == type) return &bufs[i]; return NULL; }

	int totalLength() { int length = 0; for (int i = 0; i < nb; i++) length += bufs[i].cbBuffer; return length; }

	SecBuffer& operator[](int i) { assert(i < nb); return bufs[i]; }
};

class AutoBuffers : public Buffers
{
	SecurityFunctionTable* sf;

	public:
	AutoBuffers(int nBufs, SecurityFunctionTable* pSecurityFunc): Buffers(nBufs)
	{
		sf = pSecurityFunc;
	}

	~AutoBuffers()
	{
		for (int i=0; i<nb; i++)
			FreeBuffer(i);
	}

	void FreeBuffer(int i) { if (bufs[i].pvBuffer != NULL) { sf->FreeContextBuffer(bufs[i].pvBuffer); bufs[i].pvBuffer = NULL; } }
};


namespace SSLLib
{

	SSLTransport::SSLTransport()
	{
		HMODULE hSecurity = LoadLibrary(L"Security.dll");
		INIT_SECURITY_INTERFACE pInitSecurityInterface =
			(INIT_SECURITY_INTERFACE)GetProcAddress(hSecurity, "InitSecurityInterfaceW");
		pSecurityFunc = pInitSecurityInterface();
		if (pSecurityFunc == NULL)
			throw gcnew SSLException(format("Failed to initialized security interface: %d\n", GetLastError()));

		heap = HeapCreate(0 /* default options */, 0 /* commit one page */, 0 /* grow indefinitely */);
		if (heap == NULL)
			throw gcnew SSLException(format("Failed to secure heap handle: %d\n", GetLastError()));

		sizes = new SecPkgContext_StreamSizes();
		phClientCredentials = new CredHandle();
		phContext = new CtxtHandle();
	}

	SSLTransport::!SSLTransport()
	{
		pSecurityFunc->DeleteSecurityContext(phContext);

		if (!HeapDestroy(heap))
			throw gcnew SSLException(format("Failed to destroy heap handle: %d\n", GetLastError()));
	}

	void SSLTransport::Handshake(
		String^ hostname, 
		SendCallbackDelegate^ onSendCallback, 
		FetchCallbackDelegate^ onFetchCallback,
		[System::Runtime::InteropServices::OutAttribute] array<Byte>^% extradata)
	{
		//unsigned long nPackages;
		//PSecPkgInfo pPackages;
		//SECURITY_STATUS returncode = pSecurityFunc->EnumerateSecurityPackagesW(&nPackages, &pPackages);
		//if (FAILED(returncode))
		//	throw gcnew SSLException(format("Failed to EnumerateSecurityPackages: 0x%x\n", GetLastError()));
	
		//for(unsigned long i=0; i<nPackages; i++)
		//	wprintf(L"Package %d: %s, %s\n", i, pPackages[i].Name, pPackages[i].Comment);

		//FreeContextBuffer(pPackages);

		SCHANNEL_CRED SchannelCred;
		ZeroMemory( &SchannelCred, sizeof(SchannelCred) );
		SchannelCred.dwVersion  = SCHANNEL_CRED_VERSION;
		SchannelCred.grbitEnabledProtocols = SP_PROT_TLS1;
		SchannelCred.dwFlags |= SCH_CRED_NO_DEFAULT_CREDS | SCH_CRED_MANUAL_CRED_VALIDATION;

		TimeStamp tsExpiry;

		SECURITY_STATUS returncode = pSecurityFunc->AcquireCredentialsHandle(
			NULL, // principal name
			UNISP_NAME_W, // package name
			SECPKG_CRED_OUTBOUND, // use - prepare an outgoing token
			NULL, // logon id pointer
			&SchannelCred, // package specific data
			NULL, // GetKey() fun pointer
			NULL, // GetKey() parameter
			phClientCredentials, // out - credential handle
			&tsExpiry); // out - lifetime
		if (FAILED(returncode))
			throw gcnew SSLException(format("AcquireCredentialsHandleW failed: 0x%x\n", returncode));

		const DWORD dwSSPIFlags = 
			ISC_REQ_SEQUENCE_DETECT |
			ISC_REQ_REPLAY_DETECT |
			ISC_REQ_CONFIDENTIALITY |
			ISC_RET_EXTENDED_ERROR |
			ISC_REQ_ALLOCATE_MEMORY |
			ISC_REQ_STREAM;
		DWORD dwSSPIOutFlags;


		{
			AutoBuffers bufsRequestToken(1, pSecurityFunc);
			bufsRequestToken.SetBufferToken(0, NULL, 0);

			marshal_context^ mcHostname = gcnew marshal_context();
			wchar_t* pchhostname = const_cast<wchar_t*>(mcHostname->marshal_as<const wchar_t*>(hostname));
			//
			//	requesting a TOKEN
			//
			returncode = pSecurityFunc->InitializeSecurityContextW(
				phClientCredentials,	// phCredential
				NULL,					// phContext - Context Handle is NULL - requesting a TOKEN
				pchhostname,			// pszTargetName
				dwSSPIFlags,			// fContextReq
				0,						// Reserved1
				SECURITY_NATIVE_DREP,	// TargetDataRep
				NULL,					// pInput - NULL on the first call, subsequent calls - ptr to a buffer allocated w/enough memory to hold the token
				0,						// Reserved2
				phContext,				// phNextContext - on the first call, this ptr receives new context handle. On subsequent - same as phContext up above
				&bufsRequestToken,		// pOutput - output data holder
				&dwSSPIOutFlags,		// pfContextAttr - attrs of the established ctxt
				&tsExpiry);				// ptsExpiry
			if (returncode != SEC_I_CONTINUE_NEEDED)
				throw gcnew SSLException(format("InitializeSecurityContext failed: %d. Expected follow up security context initialization.", returncode));

			CallSend(onSendCallback, (char*)bufsRequestToken[0].pvBuffer, bufsRequestToken[0].cbBuffer);
		}

		//
		// Client handshake loop
		//
		PCHAR iobuffer = (PCHAR)HeapAlloc(heap, HEAP_ZERO_MEMORY, MAX_BUFFER_SIZE);
		if (iobuffer == NULL)
			throw gcnew SSLException(format("Failed to allocate %d bytes: %d\n", MAX_BUFFER_SIZE, GetLastError()));
		DWORD cbIOBuffer = 0;
		Buffers bufsIn(2);
		AutoBuffers bufsOut(1, pSecurityFunc);

		SECURITY_STATUS sc = SEC_I_CONTINUE_NEEDED;
		while (sc == SEC_I_CONTINUE_NEEDED || sc == SEC_E_INCOMPLETE_MESSAGE || sc == SEC_I_INCOMPLETE_CREDENTIALS)
		{
			//printf("Client handshake loop: sc=%d, cbIOBuffer=%d\n", sc, cbIOBuffer);
			if (cbIOBuffer == 0 || sc == SEC_E_INCOMPLETE_MESSAGE || sc == SEC_I_CONTINUE_NEEDED)
			{
				int len = CallFetch(onFetchCallback, iobuffer + cbIOBuffer, MAX_BUFFER_SIZE - cbIOBuffer);
				if (len == 0)
				{
					printf("Server disconnected unexpectedly\n");
					sc = SEC_E_INTERNAL_ERROR;
				}
				else
					cbIOBuffer += len;
			}

			bufsIn.SetBufferToken(0, iobuffer, cbIOBuffer);
			bufsIn.SetBufferEmpty(1);

			bufsOut.SetBufferToken(0, NULL, 0);

			returncode = pSecurityFunc->InitializeSecurityContextW(
				phClientCredentials,
				phContext,
				NULL,
				dwSSPIFlags,
				0,
				SECURITY_NATIVE_DREP,
				&bufsIn,
				0,
				NULL,
				&bufsOut,
				&dwSSPIOutFlags,
				&tsExpiry);

			if (returncode == SEC_E_OK || 
				returncode == SEC_I_CONTINUE_NEEDED || 
				returncode == SEC_E_INCOMPLETE_MESSAGE ||
				FAILED(returncode) && (dwSSPIOutFlags & ISC_RET_EXTENDED_ERROR))
			{
				if (bufsOut[0].cbBuffer != 0 && bufsOut[0].pvBuffer != NULL)
				{
					CallSend(onSendCallback, (char*)bufsOut[0].pvBuffer, bufsOut[0].cbBuffer);
					bufsOut.FreeBuffer(0);
				}
			}

			if (returncode == SEC_E_INCOMPLETE_MESSAGE)
				continue;

			if (returncode == SEC_E_OK)
			{
				if (bufsIn[1].BufferType == SECBUFFER_EXTRA)
				{
//					printf("Received extra %d bytes of data with handshake data\n", bufsIn[1].cbBuffer);
					extradata = gcnew array<Byte>(bufsIn[1].cbBuffer);
					System::Runtime::InteropServices::Marshal::Copy(IntPtr(iobuffer + cbIOBuffer - bufsIn[1].cbBuffer), extradata, 0, bufsIn[1].cbBuffer);
				}
				else
					extradata = nullptr;
				break;
			}

			if (FAILED(returncode))
			{
				printf("Second InitializeSecurityContext failed : 0x%x\n", returncode);
				break;
			}

			if (returncode == SEC_I_INCOMPLETE_CREDENTIALS)
			{
				printf("Server didn't like our client certificate. Oh, well.\n");
				break;
			}

			//
			//	reuse "extra" leftovers
			//
			if (bufsIn[1].BufferType == SECBUFFER_EXTRA)
			{
				MoveMemory(iobuffer, iobuffer + (cbIOBuffer - bufsIn[1].cbBuffer), bufsIn[1].cbBuffer);
				cbIOBuffer = bufsIn[1].cbBuffer;
			}
			else
				cbIOBuffer = 0;
		}

		HeapFree(heap, 0 /* flags */, iobuffer);

		returncode = pSecurityFunc->QueryContextAttributesW(phContext, SECPKG_ATTR_STREAM_SIZES, sizes);
		if (FAILED(returncode))
			throw gcnew SSLException(format("Failed to query context attributes: 0x%x\n", returncode));
	}

	array<Byte>^ SSLTransport::Encrypt(array<Byte>^ plain_data)
	{
		pin_ptr<Byte> pinned_data = &plain_data[0];
		Byte* current_data_to_encode = pinned_data;
		int plain_length = plain_data->Length;

		char* iobuffer = (PCHAR)HeapAlloc(heap, HEAP_ZERO_MEMORY, sizes->cbMaximumMessage + sizes->cbHeader + sizes->cbTrailer);
		if (iobuffer == NULL)
			throw gcnew SSLException(format("Failed to allocate %d bytes: %d\n", plain_length, GetLastError()));
		try
		{
			array<Byte>^ outbuffer = nullptr;
			Byte* current_outbuffer = NULL;

			int remainder_length = plain_length;
			int request_length = 0;
			const int MAX_CARRY_LOAD = sizes->cbMaximumMessage - sizes->cbHeader - sizes->cbTrailer;

			while (remainder_length > 0)
			{
				request_length = min(remainder_length, MAX_CARRY_LOAD);
				remainder_length -= request_length;

				Buffers bMessage(4);
				bMessage.SetBufferStreamHeader(0, iobuffer, sizes->cbHeader);
				memcpy(iobuffer + sizes->cbHeader, current_data_to_encode, request_length);
				bMessage.SetBufferData(1, iobuffer + sizes->cbHeader, request_length);
				bMessage.SetBufferStreamTrailer(2, iobuffer + sizes->cbHeader + request_length, sizes->cbTrailer);
				bMessage.SetBufferEmpty(3);

				int returncode = pSecurityFunc->EncryptMessage(phContext, 0, &bMessage, 0);
				if (FAILED(returncode))
					throw gcnew SSLException(format("Failed to encrypt message: 0x%x\n", returncode));

				array<Byte>^ newoutbuffer = gcnew array<Byte>(bMessage.totalLength() + (outbuffer != nullptr? outbuffer->Length: 0));
				pin_ptr<Byte> pinned_newout_buffer = &newoutbuffer[0];
				if (outbuffer != nullptr)
				{
					pin_ptr<Byte> pinned_out_buffer = &outbuffer[0];
					memcpy(pinned_newout_buffer, pinned_out_buffer, outbuffer->Length);
				}
				Byte* current_outbuffer = pinned_newout_buffer + ((outbuffer != nullptr)? outbuffer->Length: 0);
				memcpy(current_outbuffer, iobuffer, bMessage.totalLength());

				outbuffer = newoutbuffer;

				current_data_to_encode += request_length;
			}

			return outbuffer;
		}
		finally
		{
			HeapFree(heap, 0 /* flags */, iobuffer);
		}
	}

	int SSLTransport::Decrypt(
			array<Byte>^ encrypted_data,  
			SendCallbackDelegate^ onSendCallback, 
			FetchCallbackDelegate^ onFetchCallback,
			SendCallbackDelegate^ onDecryptedDataCallback)
	{
		int encrypted_length = encrypted_data->Length;
//		printf("Decrypt: encrypted_length = %d\n", encrypted_length);
		if (encrypted_length > 0)
		{
			pin_ptr<Byte> pinned_data = &encrypted_data[0];

			Byte* iobuffer = (PBYTE)HeapAlloc(heap, HEAP_ZERO_MEMORY, encrypted_length);
			if (iobuffer == NULL)
				throw gcnew SSLException(format("Failed to allocate %d bytes: %d\n", encrypted_length, GetLastError()));
			try
			{
				memcpy(iobuffer, pinned_data, encrypted_length);

				while(true)
				{
					Buffers bufs(4); // buffers arrangements is dictated by MS Schannel decryption policy
					bufs.SetBufferData(0, iobuffer, encrypted_length);
					bufs.SetBufferEmpty(1);
					bufs.SetBufferEmpty(2);
					bufs.SetBufferEmpty(3);

					SECURITY_STATUS returncode = pSecurityFunc->DecryptMessage(phContext, &bufs, 0, NULL);

					if (returncode == SEC_E_INCOMPLETE_MESSAGE)
					{
						//
						//	need more data
						//
						Byte* iobuffer_more = (PBYTE)HeapAlloc(heap, HEAP_ZERO_MEMORY, MAX_BUFFER_SIZE + encrypted_length);
						if (iobuffer_more == NULL)
							throw gcnew SSLException(format("Failed to allocate %d bytes: %d\n", MAX_BUFFER_SIZE + encrypted_length, GetLastError()));
//						printf("\tBuf size is %d", MAX_BUFFER_SIZE + encrypted_length);
						int encrypted_length_more = CallFetch(onFetchCallback, iobuffer_more + encrypted_length, MAX_BUFFER_SIZE);
						if (encrypted_length_more == 0)
							throw gcnew SSLException(format("Server disconnected unexpectedly"));
//						else
//							printf("...received %d bytes of data\n", encrypted_length_more);
//						printf("\tEncrypted length = %d\n", encrypted_length);

						memcpy(iobuffer_more, iobuffer, encrypted_length);
						HeapFree(heap, 0 /* flags */, iobuffer);
						iobuffer = iobuffer_more;
						encrypted_length += encrypted_length_more;
					}
					else if (returncode == SEC_I_CONTEXT_EXPIRED)
					{
						CallSend(onSendCallback, new char[0], 0); // part of the "shutdown" protocol
						//throw gcnew SSLException(format("Session expired"));
						break;
					}
					else
					{
						if (returncode != SEC_E_OK)
							throw gcnew SSLException(format("Decryption failed: 0x%x", returncode));

						SecBuffer* pDataBuffer = bufs.First(SECBUFFER_DATA);
						if (pDataBuffer != NULL)
//						{
//							printf("\tDone decrypting %d chars\n", pDataBuffer->cbBuffer);
							CallSend(onDecryptedDataCallback, pDataBuffer->pvBuffer, pDataBuffer->cbBuffer);
//						}

						SecBuffer* pExtraBuffer = bufs.First(SECBUFFER_EXTRA);
						if (pExtraBuffer != NULL)
						{
							//
							// get extra and keep decrypting
							//
							MoveMemory(iobuffer, pExtraBuffer->pvBuffer, pExtraBuffer->cbBuffer);
							encrypted_length = pExtraBuffer->cbBuffer;
						}
						else
							break;
					}
				}
			}
			finally
			{
				HeapFree(heap, 0 /* flags */, iobuffer);
			}
		}
		return 0;
	}

	void SSLTransport::Done()
	{
	}
}

